import base64
import logging
import secrets
from typing import Optional

import click

from configs import app_config
from constants import UUID_NIL
from constants.languages import languages
from core.entities.account_role import AccountRole
from extensions.ext_database import db
from extensions.ext_redis import redis_client
from libs.helper import email as email_validate
from libs.password import hash_password, password_pattern, valid_password
from libs.regx_validator import RegxValidator
from libs.rsa import generate_key_pair
from models import Tenant
from models.account import (
    Account,
    AccountStatus,
    TenantAccountJoin,
    
)
from services.account_service import RegisterService, TenantService
from services.area_init_service import initialize_area_codes
from services.wechat_info_service import WeChatAppType, WeChatConfigService


@click.command("test")
def test():
    print(WeChatConfigService.get_wechat_tenant_info("00000000-0000-0000-0000-000000000000","wx46468a08d0d6e9e1",WeChatAppType.WECHAT_MINI_APP))

@click.command("init-area-code", help="初始化地区码")
def init_area_code():
    """
    初始化地区码
    """
    initialize_area_codes()
    click.echo(click.style("行政区划代码初始化完成！", fg="green"))
    
@click.command("reset-password", help="Reset the account password.")
@click.option("--email", prompt=True, help="Account email to reset password for")
@click.option("--new-password", prompt=True, help="New password")
@click.option("--password-confirm", prompt=True, help="Confirm new password")
def reset_password(email, new_password, password_confirm):
    """
    Reset password of owner account
    Only available in SELF_HOSTED mode
    """
    if str(new_password).strip() != str(password_confirm).strip():
        click.echo(click.style("Passwords do not match.", fg="red"))
        return

    account = db.session.query(Account).filter(Account.email == email).one_or_none()

    if not account:
        click.echo(click.style("Account not found for email: {}".format(email), fg="red"))
        return

    try:
        valid_password(new_password)
    except:
        click.echo(click.style("Invalid password. Must match {}".format(password_pattern), fg="red"))
        return

    # generate password salt
    salt = secrets.token_bytes(16)
    base64_salt = base64.b64encode(salt).decode()

    # encrypt password with salt
    password_hashed = hash_password(new_password, salt)
    base64_password_hashed = base64.b64encode(password_hashed).decode()
    account.password = base64_password_hashed
    account.password_salt = base64_salt
    db.session.commit()
    click.echo(click.style("Password reset successfully.", fg="green"))


@click.command("reset-email", help="Reset the account email.")
@click.option("--email", prompt=True, help="Current account email")
@click.option("--new-email", prompt=True, help="New email")
@click.option("--email-confirm", prompt=True, help="Confirm new email")
def reset_email(email, new_email, email_confirm):
    """
    Replace account email
    :return:
    """
    if str(new_email).strip() != str(email_confirm).strip():
        click.echo(click.style("New emails do not match.", fg="red"))
        return

    account = db.session.query(Account).filter(Account.email == email).one_or_none()

    if not account:
        click.echo(click.style("Account not found for email: {}".format(email), fg="red"))
        return

    try:
        email_validate(new_email)
    except:
        click.echo(click.style("Invalid email: {}".format(new_email), fg="red"))
        return

    account.email = new_email
    db.session.commit()
    click.echo(click.style("Email updated successfully.", fg="green"))


@click.command(
    "reset-encrypt-key-pair",
    help="Reset the asymmetric key pair of workspace for encrypt LLM credentials. "
    "After the reset, all LLM credentials will become invalid, "
    "requiring re-entry."
    "Only support SELF_HOSTED mode.",
)
@click.confirmation_option(
    prompt=click.style(
        "Are you sure you want to reset encrypt key pair? This operation cannot be rolled back!", fg="red"
    )
)
def reset_encrypt_key_pair():
    """
    Reset the encrypted key pair of workspace for encrypt LLM credentials.
    After the reset, all LLM credentials will become invalid, requiring re-entry.
    Only support SELF_HOSTED mode.
    """
    if app_config.EDITION != "SELF_HOSTED":
        click.echo(click.style("This command is only for SELF_HOSTED installations.", fg="red"))
        return

    tenants = db.session.query(Tenant).all()
    for tenant in tenants:
        if not tenant:
            click.echo(click.style("No workspaces found. Run /install first.", fg="red"))
            return

        tenant.encrypt_public_key = generate_key_pair(tenant.id)

        # db.session.query(Provider).filter(Provider.provider_type == "custom", Provider.tenant_id == tenant.id).delete()
        # db.session.query(ProviderModel).filter(ProviderModel.tenant_id == tenant.id).delete()
        db.session.commit()

        click.echo(
            click.style(
                "Congratulations! The asymmetric key pair of workspace {} has been reset.".format(tenant.id),
                fg="green",
            )
        )


@click.command("create-tenant", help="Create account and tenant.")
@click.option("--email", prompt=True, help="Tenant account email.")
@click.option("--name", prompt=True, help="Workspace name.")
@click.option("--language", prompt=True, help="Account language, default: en-US.")
def create_tenant(email: str, language: Optional[str] = None, name: Optional[str] = None):
    """
    Create tenant account
    """
    if not email:
        click.echo(click.style("Email is required.", fg="red"))
        return

    # Create account
    email = email.strip()

    if "@" not in email:
        click.echo(click.style("Invalid email address.", fg="red"))
        return

    account_name = email.split("@")[0]

    if language not in languages:
        language = "en-US"

    # Validates name encoding for non-Latin characters.
    name = name.strip().encode("utf-8").decode("utf-8") if name else None

    # generate random password
    new_password = secrets.token_urlsafe(16)

    # register account
    account = RegisterService.register(
        email=email,
        phone=None,
        name=account_name,
        password=new_password,
        language=language,
        create_workspace_required=False,
    )
    TenantService.create_owner_tenant_if_not_exist(account, name)

    click.echo(
        click.style(
            "Account and tenant created.\nAccount: {}\nPassword: {}".format(email, new_password),
            fg="green",
        )
    )


@click.command("init-tenant", help="init tenant.")
@click.option("--name", prompt=False, help="tenant name")
@click.option("--phone", prompt=False, help="Phone Number for super admin")
@click.option("--email", prompt=False, help="Email for super admin")
@click.option("--pwd", prompt=False, help="管理员密码，默认手机后六位")
def init_tenant(name:str,phone: str, pwd: str|None = None,email:str|None=None):

    if phone:
        if not RegxValidator.is_china_phone(phone) and not RegxValidator.is_malaysia_phone(
            phone
        ):
            click.echo(click.style(f"手机号不合法{phone}，请检查。", fg="red"))
            return
    # Create account
    elif email:
        if "@" not in email:
            click.echo(click.style("Invalid email address.", fg="red"))
            return
    else:
        click.echo(click.style(f"手机号和邮箱必须提供一个。", fg="red"))
        return
        # email = f"{phone}{dify_config.ACCOUNT_EMAIL_SUFFIX}"

    

    language = "en-US"
    if language not in languages:
        language = "en-US"

    # generate random password
    if pwd:
        new_password = pwd
    else:
        new_password = f"qwer{phone[-6:]}"
    # args = {
    #     "language": "en-US",
    #     "md5_salt": None,
    #     "sub_domain": None,
    #     "owner_phone": phone,
    #     "tenant_name": "SystemTenant",
    #     "password": new_password,
    #     "email": email,
    # }
    try:
        account = RegisterService.register(
            email=email,
            phone=phone,
            name='SuperAdmin',
            password=new_password,
            language=language,
            status=AccountStatus.ACTIVE,
            is_setup=True,
            create_workspace_required=False,
            force=True,
        )
        account.id = UUID_NIL
        if not name:
            name = "SystemTenant"
        tenant = TenantService.create_tenant(name=name, is_setup=True,new_tenant_id=UUID_NIL)
        account_join = TenantAccountJoin(
            tenant_id = UUID_NIL,
            account_id = UUID_NIL,
            role = AccountRole.SUPER_ADMIN.value,
            current = True,
        )
        db.session.add(account)
        db.session.add(account_join)
        
        db.session.commit()
    except Exception as e:
        db.session.rollback()
        logging.error(e)
        click.echo(click.style(f"租户创建失败，请检查。{e}", fg="red"))
        raise e
        return
    click.echo(
        click.style(
            "系统默认租户创建成功.\nAccount: Email:{}Phone:{}\n租户: {}".format(
                account.email,account.phone, tenant.name
            ),
            fg="green",
        )
    )

@click.command("upgrade-db", help="Upgrade the database")
def upgrade_db():
    click.echo("Preparing database migration...")
    lock = redis_client.lock(name="db_upgrade_lock", timeout=60)
    if lock.acquire(blocking=False):
        try:
            click.echo(click.style("Starting database migration.", fg="green"))

            # run db migration
            import flask_migrate  # type: ignore

            flask_migrate.upgrade()

            click.echo(click.style("Database migration successful!", fg="green"))

        except Exception:
            logging.exception("Failed to execute database migration")
        finally:
            lock.release()
    else:
        click.echo("Database migration skipped")


